package musicalgo;

import ai.neuralnet.ClassificationDataSet;
import ai.neuralnet.NNTrainer;
import ai.neuralnet.NeuralNetwork;
import ai.neuralnet.gradientDescent.GradientDescentTrainer;
import data.SequentialData;
import musicalgo.data.DataIO;

public class NeuralNetworkAlgorithm extends RecognitionAlgorithm
{
    protected NNTrainer trainer;

    protected NeuralNetwork bestNetwork;

    protected double trainingPercentage;

    protected int dataSetOffset;


    public NeuralNetworkAlgorithm(NNTrainer trainer)
    {
        this.trainer = trainer;
        trainingPercentage = 0.7;
        dataSetOffset = 0;
    }


    @Override
    protected void train()
    {
        ClassificationDataSet trainingSet = getTrainingSet(trainingPercentage, dataSetOffset);
        trainer.setTrainingSet(trainingSet);
        System.out.println("Training set loaded with " + trainingSet.size() + " Songs");

        ClassificationDataSet validationSet = getValidationSet(trainingPercentage, dataSetOffset);
        trainer.setValidationSet(validationSet);
        System.out.println("Validation set loaded with " + validationSet.size() + " Songs");


        System.out.println("Training Networks");
        trainer.trainNetworks();
        System.out.println("Validating Networks");
        bestNetwork = trainer.validate();

    }

    @Override
    protected void test()
    {

        System.out.println("Classifying songs ...");

        int[] predictions = getTestingSet().getPredictions(bestNetwork);

        System.out.println("Writing output");

        DataIO.writePredictions("testy.csv", predictions);
    }


    protected ClassificationDataSet getTrainingSet(double trainingPercentage, int offset)
    {
        int trainingSetSize = (int) (trainingPercentage * trainingData.size());

        ClassificationDataSet trainingSet;
        if (offset + trainingSetSize > trainingData.size())
        {
            trainingSet = getTrainingSubset(offset, trainingData.size());
            trainingSet.addAll(getTrainingSubset(0, trainingSetSize - trainingSet.size()));
        }
        else
        {
            trainingSet = getTrainingSubset(offset, offset + trainingSetSize);
        }

        return trainingSet;
    }


    /**
     * @param start inclusive, first index to include
     * @param stop  exclusive, index to stop at
     * @return subset of the training data
     */
    protected ClassificationDataSet getTrainingSubset(int start, int stop)
    {
        loadTrainingData();

        ClassificationDataSet subset = new ClassificationDataSet();
        for (int i = start + 1; i < stop; i++)
        {
            subset.addInputOutputPair(trainingData.get(i).flatData(), outputs.get(i));
        }

        return subset;
    }


    protected ClassificationDataSet getTestingSet()
    {
        loadTestingData();
        ClassificationDataSet testSet = new ClassificationDataSet();
        for (SequentialData s : testData)
        {
            testSet.addInputOutputPair(s.flatData(), null);
        }

        return testSet;
    }

    protected ClassificationDataSet getValidationSet(double trainingPercentage, int offset)
    {
        int trainingSetSize = (int) (trainingPercentage * trainingData.size());
        int validationSetSize = trainingData.size() - trainingSetSize;
        int validationOffset = (offset + trainingSetSize) % trainingData.size();

        ClassificationDataSet validationSet;
        if (validationOffset + validationSetSize > trainingData.size())
        {
            validationSet = getTrainingSubset(validationOffset, trainingData.size());
            validationSet.addAll(getTrainingSubset(0, validationSetSize - validationSet.size()));
        }
        else
        {
            validationSet = getTrainingSubset(validationOffset, validationOffset + validationSetSize);
        }

        return validationSet;
    }

    public NNTrainer getTrainer()
    {
        return trainer;
    }

    public void setTrainer(GradientDescentTrainer trainer)
    {
        this.trainer = trainer;
    }


    public void setTrainingPercentage(double trainingPercentage)
    {
        this.trainingPercentage = trainingPercentage;
    }

    public void setDataSetOffset(int dataSetOffset)
    {
        this.dataSetOffset = dataSetOffset;
    }
}
